from torch.autograd import Variable
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        # 1图像输入通道，6输出通道，5x5卷积核
        self.conv1=nn.Conv2d(1,6,5)
        self.conv2=nn.Conv2d(6,16,5)
        # 仿射操作？？线性操作y=ax+b
        self.fc1=nn.Linear(16*5*5,120)
        self.fc2=nn.Linear(120,84)
        self.fc3=nn.Linear(84,10)

    def forward(self,x):
        # 最大池化(2,2)
        x=F.max_pool2d(F.relu(self.conv1(x)),(2,2))
        # 如果size是正方形，池化可以只指定一个参数
        x=F.max_pool2d(F.relu(self.conv2(x)),2)
        x=x.view(-1,self.num_flat_features(x))
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

net=Net()
# print(net)
params=list(net.parameters())
# print(len(params))
print(params[1]) # params中存储的是所有的参数，偶数层为线上的，奇数层为点上的

inputs=Variable(torch.rand(1,1,32,32))  # Conv2d对应的输入(样本数，通道数，高度，宽度)
out=net(inputs)
print(out)
# print(net.forward(inputs)) 效果同上

# net.zero_grad()     # 清空梯度
# out.backward(torch.rand(1,10))


target=Variable(torch.arange(1,11))
criterion=nn.MSELoss()
loss=criterion(out,target)
print(loss)
loss.backward()


learning_rate=0.01
for f in net.parameters():
    f.data+=f.grad.data*learning_rate  #这两步用来更新参数

for i in range(100):
    optimizer=optim.SGD(net.parameters(),lr=0.01)
    optimizer.zero_grad()   # 清空梯度
    output=net(inputs)
    loss=criterion(output,target)   # 计算损失
    loss.backward()     # 反向传播算出梯度，backward会把该处到之前所有地方的梯度都求出来
    optimizer.step()    # 进行参数更新
